﻿namespace Cellfish.Toolbox.Tasks
{
    using System;
    using System.Threading;
    using System.Threading.Tasks;

    /// <summary>
    /// Extension methods to make wait handles awaitable as tasks.
    /// </summary>
    public static class WaitHandleExtensions
    {
        /// <summary>
        /// Converts a WaitHandle into a task.
        /// </summary>
        /// <param name="handle">The handle.</param>
        /// <returns>A task that completes when the wait handle is set.</returns>
        /// <exception cref="System.ArgumentNullException">handle</exception>
        public static Task WaitAsync(this WaitHandle handle)
        {
            if (handle == null)
            {
                throw new ArgumentNullException("handle");
            }

            if (handle.WaitOne(TimeSpan.Zero))
            {
                return Task.FromResult(true);
            }

            var tcs = new TaskCompletionSource<bool>();
            ThreadPool.RegisterWaitForSingleObject(
                handle,
                (state, isTimeout) => tcs.TrySetResult(!isTimeout),
                null,
                -1,
                true);

            return tcs.Task;
        }

        /// <summary>
        /// Converts a WaitHandle into a task.
        /// </summary>
        /// <param name="handle">The handle.</param>
        /// <param name="timeout">The timeout.</param>
        /// <returns>A task that completes when the wait handle is set.
        /// The task result will be false if the WaitHandle was not set within the
        /// given timeout, otherwise false.</returns>
        /// <exception cref="System.ArgumentNullException">handle</exception>
        public static Task<bool> WaitAsync(this WaitHandle handle, TimeSpan timeout)
        {
            if (handle == null)
            {
                throw new ArgumentNullException("handle");
            }

            if (handle.WaitOne(TimeSpan.Zero))
            {
                return Task.FromResult(true);
            }

            var tcs = new TaskCompletionSource<bool>();
            ThreadPool.RegisterWaitForSingleObject(
                handle,
                (state, isTimeout) => tcs.TrySetResult(!isTimeout),
                null,
                timeout,
                true);

            return tcs.Task;
        }
    }
}
